from matrix import Matrix
from vector import Vector


class LinearSystem:

    def __init__(self, A: Matrix, b: Vector):
        self._m = A.row_num()
        self._n = A.col_num()
        self.Ab = [Vector(A.row_vector(i)._values + [b[i]]) for i in range(self._m)]

    def __repr__(self):
        return f"LinearSystem({self.Ab})"

    def fancy_print(self):

        for i in range(self._m):
            print(" ".join(str(self.Ab[i][j]) for j in range(self._n)), end=" ")
            print("|", self.Ab[i][-1])

    def max_row(self, start, end):
        best, ret = abs(self.Ab[start][start]), start
        for i in range(start + 1, end):
            if best < abs(self.Ab[i][start]):
                best, ret = abs(self.Ab[i][start]), i
        return ret

    def forward(self):
        for i in range(self._m):
            ret = self.max_row(i, self._m)
            self.Ab[ret], self.Ab[i] = self.Ab[i], self.Ab[ret]

            self.Ab[i] = self.Ab[i] / self.Ab[i][i]
            for j in range(i + 1, self._m):
                # self.Ab[j] -= self.Ab[i] * (self.Ab[j][i] / self.Ab[i][i])
                self.Ab[j] -= self.Ab[i] * (self.Ab[j][i]/self.Ab[i][i])
        # for i in range(self._m):
        #     for j in range(i+1,self._m):
        #         self.Ab[j] = self.Ab[j] - self.Ab[i] * (self.Ab[j][i]/self.Ab[i][i])

    def backward(self):
        for i in range(self._m - 1, -1, -1):
            for j in range(i - 1, -1, -1):
                # self.Ab[j] -= self.Ab[i] * (self.Ab[j][i] / self.Ab[i][i])
                self.Ab[j] = self.Ab[j] - self.Ab[j][i] * self.Ab[i]



    def gauss_jordan_elimination(self):
        self.forward()
        self.backward()

    __str__ = __repr__


if __name__ == '__main__':
    # m1 = Matrix([[1, 2, 3],
    #              [4, 5, 6],
    #              [7, 8, 9]])
    # b = [10, 20, 30]
    #
    # linear01 = LinearSystem(m1, b)
    # print(linear01)
    # linear01.gauss_jordan_elimination()
    # print(linear01)

    A = Matrix([
        [1,3,1], [2,1,1], [2, 2, 1],
    ])
    b = Vector([11, 8, 10])
    Ab = LinearSystem(A, b)

    print("------------")
    # print(Ab)


    Ab.gauss_jordan_elimination()
    # print(Ab)
    Ab.fancy_print()
